
#include "histogram_tiling.h"
#include "register/op_def_registry.h"
#include "tiling_common.h"
#include <cstdint>
namespace optiling
{
static ge::graphStatus TilingFunc(gert::TilingContext *context)
{

    HistogramTilingData tiling;
    context->SetBlockDim(1);
    auto xshape = context->GetInputShape(0)->GetStorageShape();
    auto size = xshape.GetShapeSize();
    tiling.set_size(size);

    auto att_bins = context->GetAttrs()->GetInt(0);
    auto att_min = context->GetAttrs()->GetFloat(1);
    auto att_max = context->GetAttrs()->GetFloat(2);
    tiling.set_bins(*att_bins);
    tiling.set_min(*att_min);
    tiling.set_max(*att_max);

    auto dtype = context->GetInputDesc(0)->GetDataType();
    auto dtypeSize = GetSizeByDataType(dtype);
    int64_t alignSize = 256 / dtypeSize;

    COMMON_TILING_GET_PLATFORM;
    COMMON_TILING_GET_UB_SIZE;
    double tileLenthCount = 2.0 + 2.0 + (1.0 / double(dtypeSize)) + (2 / double(alignSize));
    if (dtype == ge::DataType::DT_FLOAT16)
    {
        tileLenthCount = 2.0 + 2 * tileLenthCount;
    }
    else if (dtype == ge::DataType::DT_INT32)
    {
        tileLenthCount = 1 + tileLenthCount;
    }
    auto extra = 256;
    if ((*att_min) == 0 && (*att_max) == 0)
    {
        extra += 32 * 4;
    }
    auto tileLength = int64_t(double(ub_size - extra) / tileLenthCount) / 256 * 256 / dtypeSize;
    if (size < tileLength)
    {
        tileLength = ALIGN_TO(size, alignSize);
    }
    tiling.set_tileLength(tileLength);
    auto finalLength = size % tileLength;
    if (finalLength == 0)
    {
        finalLength = size;
    }
    tiling.set_finalLength(finalLength);

    tiling_zero_workspace(context);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

    std::cout << "bins = " << tiling.get_bins() << std::endl;
    std::cout << "min = " << tiling.get_min() << std::endl;
    std::cout << "max = " << tiling.get_max() << std::endl;
    std::cout << "tileLenth = " << tiling.get_tileLength() << std::endl;
    std::cout << "finalLenth = " << tiling.get_finalLength() << std::endl;
    std::cout << "size = " << tiling.get_size() << std::endl;

    return ge::GRAPH_SUCCESS;
}
} // namespace optiling

namespace ge
{
static ge::graphStatus InferShape(gert::InferShapeContext *context)
{
    auto att_bins = context->GetAttrs()->GetInt(0);
    gert::Shape *y_shape = context->GetOutputShape(0);
    y_shape->SetDim(0, *att_bins);
    y_shape->SetDimNum(1);
    return GRAPH_SUCCESS;
}
} // namespace ge

namespace ops
{
class Histogram : public OpDef
{
  public:
    explicit Histogram(const char *name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        // this->Input("x")
        //     .ParamType(REQUIRED)
        //     .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
        //     .Format({ge::FORMAT_ND, ge::FORMAT_ND})
        //     .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        // this->Output("y")
        //     .ParamType(REQUIRED)
        //     .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
        //     .Format({ge::FORMAT_ND, ge::FORMAT_ND})
        //     .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("bins").AttrType(OPTIONAL).Int(100);
        this->Attr("min").AttrType(OPTIONAL).Float(0);
        this->Attr("max").AttrType(OPTIONAL).Float(0);

        this->SetInferShape(ge::InferShape);

        this->AICore().SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");
    }
};

OP_ADD(Histogram);
} // namespace ops
